from model import UNet
from data_process import Data_Loader
import torch
import torch.nn as nn
import numpy as np
from PIL import Image
import os

# 确保保存图片的文件夹存在
save_dir = "result"
if not os.path.exists(save_dir):
    os.makedirs(save_dir)

# 定义计算指标的函数
def calculate_metrics(pred, label):
    pred = pred.cpu().numpy().flatten()
    label = label.cpu().numpy().flatten()
    
    TP = ((pred == 1) & (label == 1)).sum()
    TN = ((pred == 0) & (label == 0)).sum()
    FP = ((pred == 1) & (label == 0)).sum()
    FN = ((pred == 0) & (label == 1)).sum()
    
    ACC = (TP + TN) / (TP + TN + FP + FN)
    PPV = TP / (TP + FP) if (TP + FP) != 0 else 0
    TPR = TP / (TP + FN) if (TP + FN) != 0 else 0
    F1 = 2 * TP / (2 * TP + FP + FN) if (2 * TP + FP + FN) != 0 else 0
    JS = TP / (TP + FP + FN) if (TP + FP + FN) != 0 else 0
    
    return ACC, PPV, TPR, F1, JS

# 加载测试集
testset = Data_Loader("processed_data/test/images")
test_loader = torch.utils.data.DataLoader(dataset=testset, batch_size=1)

# 加载网络和最优模型参数
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net = UNet(in_channels=1, num_classes=1)
net.to(device)
net.load_state_dict(torch.load('UNet.pth'))  # 加载最优模型参数
net.eval()

# 测试并计算指标
total_ACC, total_PPV, total_TPR, total_F1, total_JS = 0, 0, 0, 0, 0
total = 0
with torch.no_grad():
    for i, (test_image, test_label) in enumerate(test_loader):
        outputs = net(test_image.to(device))
        outputs[outputs >= 0] = 1
        outputs[outputs < 0] = 0
        
        ACC, PPV, TPR, F1, JS = calculate_metrics(outputs, test_label.to(device))
        total_ACC += ACC
        total_PPV += PPV
        total_TPR += TPR
        total_F1 += F1
        total_JS += JS
        total += 1
        
        # 打印每张图片的评价结果
        print(f"Image {i+1} - ACC: {ACC:.4f}, PPV: {PPV:.4f}, TPR: {TPR:.4f}, F1: {F1:.4f}, JS: {JS:.4f}")
        
        # 保存预测的图片
        pred_image = outputs.squeeze().cpu().numpy() * 255  # 将预测结果转换为0-255的像素值
        pred_image = pred_image.astype(np.uint8)
        image_path = os.path.join(save_dir, f"predicted_image_{i+1}_Unet.png")
        Image.fromarray(pred_image).save(image_path)

print(f"Average ACC: {total_ACC / total:.4f}")
print(f"Average PPV: {total_PPV / total:.4f}")
print(f"Average TPR: {total_TPR / total:.4f}")
print(f"Average F1: {total_F1 / total:.4f}")
print(f"Average JS: {total_JS / total:.4f}")